import os
import requests
import time
from fastapi import APIRouter

from ai_img.task import query_task_status

os.environ["DASHSCOPE_API_KEY"] = "sk-989ed32f49284be49d87158b4949b003"
DASHSCOPE_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"
imageedit_app = APIRouter()


@imageedit_app.get("/imageedit")
def generate_image(param: str):
    # 设置 API 密钥
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("环境变量 DASHSCOPE_API_KEY 未设置")

    headers = {
            "X-DashScope-Async": "enable",
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "stylization_all",
            "prompt": param,
            "base_image_url": "http://wanx.alicdn.com/material/20250318/stylization_all_1.jpeg"
        },
        "parameters": {
            "n": 1
        }
    }

    response = requests.post(DASHSCOPE_URL, headers=headers, json=data)

    # 检查响应状态码
    if response.status_code == 200:
        print("请求成功，响应数据：")
        response_data = response.json()
        print(response_data)

        # 获取任务 ID
        task_id = response_data['output']['task_id']
        print(f"任务 ID: {task_id}")

        # 定时查询任务状态
        while True:
            status_response = query_task_status(task_id,api_key)
            status_data = status_response.json()

            # 检查任务状态
            if status_data['output']['task_status'] == 'SUCCEEDED':
                print("任务完成")
                result_url = status_data['output']['results'][0]['url']
                print(f"生成的图片 URL: {result_url}")
                return {'img_url': result_url}

            print("任务仍在处理中，等待 3 秒后重试...")
            time.sleep(3)
    else:
        print(f"请求失败，状态码：{response.status_code}")
        print(f"错误信息：{response.text}")

